from pathlib import Path
from typing import List, Tuple, Dict, Optional
import numpy as np
import time
from tqdm import tqdm
import logging
from utils import setup_logging, save_model

class HMMTrainer:
    def __init__(self, model, logger: Optional[logging.Logger] = None):
        self.model = model
        self.logger = logger or setup_logging()
        self.data = None
        
    def load_data(self, file_path: str) -> List[List[Tuple[str, str]]]:
        """加载并预处理数据"""
        try:
            encoding = 'gbk'
            with open(file_path, 'r', encoding=encoding) as file:
                content = file.read()
            self.logger.info(f"文件格式：{encoding}")
            
            tag_list = content.strip().split('\n\n')
            processed_data = []
            invalid_count = 0
            
            for sentence in tqdm(tag_list, desc="Processing data"):
                word_tags = []
                for line in sentence.split('\n'):
                    parts = line.strip().split(' ')
                    if len(parts) == 2:
                        word_tags.append(tuple(parts))
                    else:
                        invalid_count += 1
                
                if word_tags:
                    processed_data.append(word_tags)
            
            self.data = processed_data
            return processed_data
            
        except UnicodeDecodeError as e:
            self.logger.error(f"数据加载失败: 编码错误 - {str(e)}")
            # 尝试使用其他编码
            alternative_encodings = ['utf-8', 'utf-8-sig', 'gb18030','gb2312', 'big5', 'ascii']
            for alt_encoding in alternative_encodings:
                try:
                    with open(file_path, 'r', encoding=alt_encoding) as file:
                        content = file.read()
                    self.logger.info(f"文件格式：{alt_encoding}")
                    # 处理数据
                    tag_list = content.strip().split('\n\n')
                    processed_data = []
                    for sentence in tqdm(tag_list, desc="Processing data"):
                        word_tags = []
                        for line in sentence.split('\n'):
                            parts = line.strip().split(' ')
                            if len(parts) == 2:
                                word_tags.append(tuple(parts))
                        if word_tags:
                            processed_data.append(word_tags)
                    self.data = processed_data
                    return processed_data
                except Exception as e:
                    self.logger.error(f"尝试使用编码 {alt_encoding} 失败: {str(e)}")
            return []
        except Exception as e:
            self.logger.error(f"数据加载失败: {str(e)}")
            return []

    def train(self, smooth_factor: float = 0.1) -> bool:
        """训练模型"""
        if not self.data:
            self.logger.error("未加载训练数据")
            return False
            
        try:
            self.model.train(self.data, smooth_factor=smooth_factor)
            return True
        except Exception as e:
            self.logger.error(f"模型训练失败: {str(e)}")
            return False

    def evaluate(self, test_data: Optional[List[List[Tuple[str, str]]]] = None) -> Dict:
        """评估模型性能"""
        data = test_data if test_data is not None else self.data
        data = [data[i] for i in range(0, len(data), 500)]  # 采样评估
    
        if not data:
            self.logger.error("无可用的评估数据")
            return {}
            
        total_correct = 0
        total_tags = 0
        sequence_probabilities = []
        
        for sequence in tqdm(data, desc="Evaluating"):
            words = [word for word, _ in sequence]
            true_tags = [tag for _, tag in sequence]
            
            pred_tags, probability = self.model.viterbi(words)
            
            correct = sum(1 for t, p in zip(true_tags, pred_tags) if t == p)
            total_correct += correct
            total_tags += len(true_tags)
            sequence_probabilities.append(probability)
        
        return {
            'accuracy': total_correct / total_tags if total_tags > 0 else 0.0,
            'average_probability': np.mean(sequence_probabilities),
            'total_tags': total_tags,
            'correct_tags': total_correct
        }

def main():
    """训练入口"""
    from HiddenMarkovModel import HiddenMarkovModel
    
    logger = setup_logging()
    
    try:
        logger.info("=" * 50)
        logger.info("HMM模型训练与优化开始")
        logger.info("=" * 50)
        
        # 初始化
        model = HiddenMarkovModel()
        trainer = HMMTrainer(model, logger)
        
        # 加载数据
        data_path = Path(r"D:\NLP-program\基于HMM的三位一体字标注汉语词法分析\flagged\pku_stn_training.txt")
        trainer.load_data(str(data_path))
        
        # 参数搜索
        smooth_factors = [0.0001,0.001, 0.01, 0.1]
        best_score = 0
        best_smooth_factor = None
        results = []
        
        start_time = time.time()
        
        for i, smooth_factor in enumerate(smooth_factors, 1):
            logger.info(f"\n[{i}/{len(smooth_factors)}] 测试平滑因子: {smooth_factor}")
            
            logger.info(f"当前平滑因子: {smooth_factor}")
            
            if trainer.train(smooth_factor=smooth_factor):
                result = trainer.evaluate()
                accuracy = result['accuracy']
                results.append((smooth_factor, accuracy))
                
                logger.info(f"准确率: {accuracy:.4f}")
                
                if accuracy > best_score:
                    best_score = accuracy
                    best_smooth_factor = smooth_factor
        
        # 输出结果
        total_time = time.time() - start_time
        logger.info(f"\n总用时: {total_time:.2f}秒")
        
        for smooth_factor, accuracy in results:
            logger.info(f"平滑因子: {smooth_factor:6.4f} -> 准确率: {accuracy:.4f}")
        
        if best_smooth_factor is not None:
            logger.info(f"\n最优平滑因子: {best_smooth_factor}")
            logger.info(f"最优准确率: {best_score:.4f}")
            
            trainer.train(smooth_factor=best_smooth_factor)
            save_model(trainer.model, r"D:\NLP-program\基于HMM的三位一体字标注汉语词法分析\flagged\hmm_model_best.pkl", logger)
            
    except Exception as e:
        logger.error(f"训练过程中发生错误: {str(e)}", exc_info=True)
    finally:
        logger.info("=" * 50)
        logger.info("训练过程结束")
        logger.info("=" * 50)

if __name__ == "__main__":
    main()